/*
* Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
* Description: Hccl One-Sided Conn
* Author: yanzhenghang
* Create: 2024-09-06
*/
#include "hccl_one_sided_conn.h"
#include "sal_pub.h"

namespace hccl {
using namespace std;

HcclOneSidedConn::HcclOneSidedConn(const HcclNetDevCtx &netDevCtx, const HcclRankLinkInfo &localRankInfo,
    const HcclRankLinkInfo &remoteRankInfo, std::unique_ptr<HcclSocketManager> &socketManager,
    LocalRdmaRmaBufferMgr &localRdmaRmaBufferMgr, std::unique_ptr<NotifyPool> &notifyPool,
    LocalIpcRmaBufferMgr &localIpcRmaBufferMgr, const HcclDispatcher &dispatcher,
    const std::string &commIdentifier, const bool &useRdma)
    : localRankInfo_(localRankInfo), socketManager_(socketManager), localRdmaRmaBufferMgr_(localRdmaRmaBufferMgr),
      localIpcRmaBufferMgr_(localIpcRmaBufferMgr), notifyPool_(notifyPool)
{
    netDevCtx_ = netDevCtx;
    remoteRankInfo_ = remoteRankInfo;
    dispatcher_ = dispatcher;
    commIdentifier_ = commIdentifier;
    useRdma_ = useRdma;
    if (useRdma) {
        transportMemPtr_ = std::make_unique<TransportRoceMem>(notifyPool, netDevCtx,
            dispatcher, localRankInfo.userRank, remoteRankInfo.userRank);
    } else {
        transportMemPtr_ = std::make_unique<TransportIpcMem>(notifyPool, netDevCtx,
            dispatcher, localRankInfo.userRank, remoteRankInfo.userRank);
    }
    CHK_SMART_PTR_RET_NULL(transportMemPtr_);
}

HcclOneSidedConn::~HcclOneSidedConn()
{
}

HcclResult HcclOneSidedConn::Connect(s32 timeoutSec)
{
    // 创建socket用于交换数据
    std::string newTag;
    if (localRankInfo_.userRank < remoteRankInfo_.userRank) {
        // 本端为SERVER，对端为CLIENT
        newTag = string(localRankInfo_.ip.GetReadableIP()) + "_" + to_string(localRankInfo_.port) + "_" +
            string(remoteRankInfo_.ip.GetReadableIP()) + "_" + to_string(remoteRankInfo_.port) + "_" + commIdentifier_;
    } else {
        newTag = string(remoteRankInfo_.ip.GetReadableIP()) + "_" + to_string(remoteRankInfo_.port) + "_" +
            string(localRankInfo_.ip.GetReadableIP()) + "_" + to_string(localRankInfo_.port) + "_" + commIdentifier_;
    }
    HCCL_DEBUG("[HcclOneSidedConn][Connect]socket tag:%s", newTag.c_str());
    std::vector<std::shared_ptr<HcclSocket>> connectSockets;
    CHK_RET(socketManager_->CreateSingleLinkSocket(newTag, netDevCtx_, remoteRankInfo_, connectSockets, true, true));
    socket_ = connectSockets[0];

    if (useRdma_) {
        // 创建socket用于QP建链
        newTag += "_QP";
        CHK_RET(socketManager_->CreateSingleLinkSocket(newTag, netDevCtx_, remoteRankInfo_, connectSockets, true, true));
        CHK_RET(transportMemPtr_->SetSocket(connectSockets[0]));
        // Transport建链：notify资源创建+QP建链
        CHK_RET(transportMemPtr_->Connect(timeoutSec));
    }
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedConn::ExchangeIpcProcessInfo(const ProcessInfo &localProcess, ProcessInfo &remoteProcess)
{
    HCCL_DEBUG("[HcclOneSidedConn][ExchangeIpcProcessInfo]localRank[%u] exchange process info", localRankInfo_.userRank);
    if (socket_->GetLocalRole() == HcclSocketRole::SOCKET_ROLE_CLIENT) {
        // 先收后发
        CHK_RET(socket_->Recv(&remoteProcess, sizeof(ProcessInfo)));
        CHK_RET(socket_->Send(&localProcess, sizeof(ProcessInfo)));
    } else {
        // 先发后收
        CHK_RET(socket_->Send(&localProcess, sizeof(ProcessInfo)));
        CHK_RET(socket_->Recv(&remoteProcess, sizeof(ProcessInfo)));
    }
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedConn::ExchangeMemDesc(const HcclMemDescs &localMemDescs, HcclMemDescs &remoteMemDescs, u32 &actualNumOfRemote)
{
    u32 localRankId = localRankInfo_.userRank;
    HCCL_DEBUG("[HcclOneSidedConn][ExchangeMemDesc]localRank[%u] exchange memDesc begin", localRankId);

    if (socket_->GetLocalRole() == HcclSocketRole::SOCKET_ROLE_CLIENT) {
        // 先收后发
        CHK_RET(ReceiveRemoteMemDesc(remoteMemDescs, actualNumOfRemote));
        CHK_RET(SendLocalMemDesc(localMemDescs));
    } else {
        // 先发后收
        CHK_RET(SendLocalMemDesc(localMemDescs));
        CHK_RET(ReceiveRemoteMemDesc(remoteMemDescs, actualNumOfRemote));
    }

    // 校验remoteDescs中的remoteRankId和conn对象中保存的localRankId是否一样
    for (u32 i = 0; i < actualNumOfRemote; i++) {
        CHK_PTR_NULL((remoteMemDescs.array) + i);
        MemDescData* ptr = reinterpret_cast<MemDescData*>(remoteMemDescs.array[i].desc);
        u32 tempRankId = ptr->remoteRankId;
        HCCL_DEBUG("[HcclOneSidedConn][ExchangeMemDesc]tempRankId:%u, localRankInfo_.userRank:%u", tempRankId, localRankId);
        if (tempRankId != localRankId) {
            HCCL_ERROR("[HcclOneSideConn][ExchangeMemDesc]localRank[%u] receive remoteMemDesc from wrong localRank[%u], "\
                "connection is for localRank[%u]", localRankId, tempRankId, localRankId);
            return HCCL_E_INTERNAL;
        }
    }
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedConn::SendLocalMemDesc(const HcclMemDescs &localMemDescs)
{
    HcclResult ret = socket_->Send(&localMemDescs.arrayLength, sizeof(u32));
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("errNo[0x%016llx] localRank[%u] send localMemDesc.arrayLength to remote "\
            "failed, ret[%u]", HCCL_ERROR_CODE(ret), localRankInfo_.userRank, ret), ret);
    HCCL_DEBUG("send localMemDescs.arrayLength:%u", localMemDescs.arrayLength);

    if (localMemDescs.arrayLength == 0) {
        HCCL_INFO("localMemDescs.arrayLength[%u], no need to send data", localMemDescs.arrayLength);
    } else {
        HCCL_DEBUG("send descSize:%u", localMemDescs.arrayLength * sizeof(HcclMemDesc));
        ret = socket_->Send(localMemDescs.array, localMemDescs.arrayLength * sizeof(HcclMemDesc));
        CHK_PRT_RET(ret != HCCL_SUCCESS,
            HCCL_ERROR("errNo[0x%016llx] localRank[%u] send localMemDesc to remote "\
                "failed, ret[%u]", HCCL_ERROR_CODE(ret), localRankInfo_.userRank, ret), ret);
    }
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedConn::ReceiveRemoteMemDesc(HcclMemDescs &remoteMemDescs, u32 &actualNumOfRemote)
{
    HcclResult ret = socket_->Recv(&actualNumOfRemote, sizeof(u32));
    remoteMemDescs.arrayLength = actualNumOfRemote;
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("errNo[0x%016llx] localRank[%u] receive actualNumOfRemote to remote "\
            "failed, ret[%u]", HCCL_ERROR_CODE(ret), localRankInfo_.userRank, ret), ret);
    HCCL_DEBUG("receive actualNumOfRemote:%u", actualNumOfRemote);
    if (actualNumOfRemote == 0) {
        HCCL_INFO("actualNumOfRemote[%u], no need to receive data", actualNumOfRemote);
    } else {
        HCCL_DEBUG("receive descSize:%u", actualNumOfRemote * sizeof(HcclMemDesc));
        ret = socket_->Recv(remoteMemDescs.array, actualNumOfRemote * sizeof(HcclMemDesc));
        CHK_PRT_RET(ret != HCCL_SUCCESS,
            HCCL_ERROR("errNo[0x%016llx] localRank[%u] receive remoteMemDesc from remote "\
                "failed, ret[%u]", HCCL_ERROR_CODE(ret), localRankInfo_.userRank, ret), ret);
    }
    return HCCL_SUCCESS;
}

void HcclOneSidedConn::EnableMemAccessRdma(const HcclMemDesc &remoteMemDesc, HcclMem &remoteMem)
{
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(remoteMemDesc.desc);
    std::string tempDesc = ptr->CopyToString();

    std::shared_ptr<RemoteRdmaRmaBuffer> tempRemoteBufferPtr = make_shared<RemoteRdmaRmaBuffer>();
    EXCEPTION_THROW_IF_ERR(tempRemoteBufferPtr->Deserialize(tempDesc),
        "[HcclOneSideConn][EnableMemAccessRdma]RemoteBuffer Deserialize failed.");
    BufferKey<uintptr_t, u64> tempKey(
        reinterpret_cast<uintptr_t>(tempRemoteBufferPtr->GetAddr()), tempRemoteBufferPtr->GetSize()
    );
    auto resultPair = remoteRdmaRmaBufferMgr_.Add(tempKey, tempRemoteBufferPtr);
    if (resultPair.first == remoteRdmaRmaBufferMgr_.End()) {
        // 输入key是表中某一个最相近key的交集、子集。返回空迭代器
        HCCL_ERROR("[HcclOneSideConn][EnableMemAccessRdma]The memory that is expected to enable"\
            " overlaps with the memory that has been enabled, please check params");
        throw logic_error("[HcclOneSideConn][EnableMemAccessRdma]Enable memory access failed.");
    }

    // 已使能：输入key是表中某一最相近key的全集。 返回添加该key的迭代器，及false
    // 未使能：输入key是表中某一最相近key的空集。 返回添加成功的迭代器，及true
    std::string logInfo = resultPair.second ? "Enable memory access success!"
                        : "Memory is already enabled, just increase the reference count.";
    HCCL_INFO("[HcclOneSideConn][EnableMemAccessRdma]:%s", logInfo.c_str());
    // 填充出参HcclMem信息
    remoteMem.addr = tempRemoteBufferPtr->GetAddr();
    remoteMem.size = tempRemoteBufferPtr->GetSize();
    remoteMem.type = tempRemoteBufferPtr->GetMemType();
}

void HcclOneSidedConn::EnableMemAccessIpc(const HcclMemDesc &remoteMemDesc, HcclMem &remoteMem)
{
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(remoteMemDesc.desc);
    std::string tempDesc = ptr->CopyToString();

    // 创建远程缓冲区，并进行反序列化、打开
    std::shared_ptr<RemoteIpcRmaBuffer> tempRemoteBufferPtr = make_shared<RemoteIpcRmaBuffer>(netDevCtx_);
    EXCEPTION_THROW_IF_ERR(tempRemoteBufferPtr->Deserialize(tempDesc),
        "[HcclOneSideConn][EnableMemAccessIpc]RemoteBuffer Deserialize failed.");
    EXCEPTION_THROW_IF_ERR(tempRemoteBufferPtr->Open(), "[HcclOneSideConn][EnableMemAccessIpc]RemoteBuffer Open failed.");
    BufferKey<uintptr_t, u64> tempKey(
        reinterpret_cast<uintptr_t>(tempRemoteBufferPtr->GetAddr()), tempRemoteBufferPtr->GetSize()
    );

    auto resultPair = remoteIpcRmaBufferMgr_.Add(tempKey, tempRemoteBufferPtr);
    if (resultPair.first == remoteIpcRmaBufferMgr_.End()) {
        // 输入key与已有的内存重叠
        HCCL_ERROR("[HcclOneSideConn][EnableMemAccessIpc]The memory that is expected to enable"\
            " overlaps with the memory that has been enabled, please check params");
        throw logic_error("[HcclOneSideConn][EnableMemAccessIpc]Enable memory access failed.");
    }

    // 已使能：输入key是表中某一最相近key的全集。 返回添加该key的迭代器，及false
    // 未使能：输入key是表中某一最相近key的空集。 返回添加成功的迭代器，及true
    std::string logInfo = resultPair.second ? "Enable memory access success!"
                                            : "Memory is already enabled, just increase the reference count.";
    HCCL_INFO("[HcclOneSideConn][EnableMemAccessIpc]:%s", logInfo.c_str());

    // 填充出参HcclMem信息
    remoteMem.addr = tempRemoteBufferPtr->GetAddr();
    remoteMem.size = tempRemoteBufferPtr->GetSize();
    remoteMem.type = tempRemoteBufferPtr->GetMemType();
}

void HcclOneSidedConn::EnableMemAccess(const HcclMemDesc &remoteMemDesc, HcclMem &remoteMem)
{
    if (useRdma_) {
        EnableMemAccessRdma(remoteMemDesc, remoteMem);
    } else {
        EnableMemAccessIpc(remoteMemDesc, remoteMem);
    }
}

void HcclOneSidedConn::DisableMemAccessRdma(const HcclMemDesc &remoteMemDesc)
{
    // 内存去使能管理
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(remoteMemDesc.desc);
    std::string tempDesc = ptr->CopyToString();

    RemoteRdmaRmaBuffer tempRemoteBuffer;
    EXCEPTION_THROW_IF_ERR(tempRemoteBuffer.Deserialize(tempDesc),
        "[HcclOneSideConn][DisableMemAccessRdma]RemoteBuffer Deserialize failed.");
    BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(tempRemoteBuffer.GetAddr()), tempRemoteBuffer.GetSize());
    if (remoteRdmaRmaBufferMgr_.Del(tempKey)) {
        // 删除成功：输入key是表中某一最相近key的全集，计数-1后为0，返回true
        HCCL_INFO("[HcclOneSidedService][DisableMemAccessRdma]Memory reference count is 0, disable memory access.");
    } else {
        // 删除失败：输入key是表中某一最相近key的全集，计数不为0（存在其他remoteRank使用），返回false
        HCCL_INFO("[HcclOneSidedService][DisableMemAccessRdma]Memory reference count is larger than 0"\
        "(used by other RemoteRank), do not disable memory.");
    }
}

void HcclOneSidedConn::DisableMemAccessIpc(const HcclMemDesc &remoteMemDesc)
{
    // 内存去使能管理
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(remoteMemDesc.desc);
    std::string tempDesc = ptr->CopyToString();

    std::shared_ptr<RemoteIpcRmaBuffer> tempRemoteBuffer = make_shared<RemoteIpcRmaBuffer>(netDevCtx_);
    EXCEPTION_THROW_IF_ERR(tempRemoteBuffer->Deserialize(tempDesc),
        "[HcclOneSideConn][DisableMemAccessIpc]RemoteBuffer Deserialize failed.");
    EXCEPTION_THROW_IF_ERR(tempRemoteBuffer->Close(),
        "[HcclOneSideConn][DisableMemAccessIpc]RemoteBuffer Close failed.");
    BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(tempRemoteBuffer->GetAddr()), tempRemoteBuffer->GetSize());
    if (remoteIpcRmaBufferMgr_.Del(tempKey)) {
        // 删除成功：输入key是表中某一最相近key的全集，计数-1后为0，返回true
        HCCL_INFO("[HcclOneSidedService][DisableMemAccessIpc]Memory reference count is 0, disable memory access.");
    } else {
        // 删除失败：输入key是表中某一最相近key的全集，计数不为0（存在其他remoteRank使用），返回false
        HCCL_INFO("[HcclOneSidedService][DisableMemAccessIpc]Memory reference count is larger than 0"\
        "(used by other RemoteRank), do not disable memory.");
    }
}

void HcclOneSidedConn::DisableMemAccess(const HcclMemDesc &remoteMemDesc)
{
    if (useRdma_) {
        DisableMemAccessRdma(remoteMemDesc);
    } else {
        DisableMemAccessIpc(remoteMemDesc);
    }
}

void HcclOneSidedConn::BatchWrite(const HcclOneSideOpDesc* oneSideDescs, u32 descNum, const rtStream_t& stream)
{
    HcclOneSideOpDesc* descPtr = const_cast<HcclOneSideOpDesc*>(oneSideDescs);
    RmaBufferSlice localRmaBufferSlice{};
    RmaBufferSlice remoteRmaBufferSlice{};
    for (u32 i = 0; i < descNum; i++) {
        if (descPtr[i].count == 0) {
            HCCL_WARNING("[HcclOneSidedConn][BatchWrite] Desc item[%u] count is 0.", i);
        }
        EXCEPTION_THROW_IF_ERR(FillRmaBufferSlice(descPtr + i, localRmaBufferSlice, remoteRmaBufferSlice),
            "[HcclOneSidedConn][BatchWrite] Get rmaBufferSlice failed!");
        EXCEPTION_THROW_IF_ERR(transportMemPtr_->Write(remoteRmaBufferSlice, localRmaBufferSlice, stream),
            "[HcclOneSidedConn][BatchWrite] transportMem Write failed.");
    }
    EXCEPTION_THROW_IF_ERR(transportMemPtr_->AddOpFence(stream), "[HcclOneSidedConn][BatchWrite] AddOpFence failed.");
}

void HcclOneSidedConn::BatchRead(const HcclOneSideOpDesc* oneSideDescs, u32 descNum, const rtStream_t& stream)
{
    HcclOneSideOpDesc* descPtr = const_cast<HcclOneSideOpDesc*>(oneSideDescs);
    RmaBufferSlice localRmaBufferSlice{};
    RmaBufferSlice remoteRmaBufferSlice{};
    for (u32 i = 0; i < descNum; i++) {
        if (descPtr[i].count == 0) {
            HCCL_WARNING("[HcclOneSidedConn][BatchRead] Desc item[%u] count is 0.", i);
        }
        EXCEPTION_THROW_IF_ERR(FillRmaBufferSlice(descPtr + i, localRmaBufferSlice, remoteRmaBufferSlice),
            "[HcclOneSidedConn][BatchRead] Get rmaBufferSlice failed!");
        EXCEPTION_THROW_IF_ERR(transportMemPtr_->Read(localRmaBufferSlice, remoteRmaBufferSlice, stream),
            "[HcclOneSidedConn][BatchRead] transportMem Read failed.");
    }
    EXCEPTION_THROW_IF_ERR(transportMemPtr_->AddOpFence(stream), "[HcclOneSidedConn][BatchRead] AddOpFence failed.");
}

HcclResult HcclOneSidedConn::FillRmaBufferSlice(const HcclOneSideOpDesc* desc,
    RmaBufferSlice& localRmaBufferSlice, RmaBufferSlice& remoteRmaBufferSlice)
{
    if (useRdma_) {
        EXCEPTION_THROW_IF_ERR(FillRdmaRmaBufferSlice(desc, localRmaBufferSlice, remoteRmaBufferSlice),
            "[HcclOneSidedConn][FillRdmaRmaBufferSlice] Get rmaBufferSlice failed!");
    } else {
        EXCEPTION_THROW_IF_ERR(FillIpcRmaBufferSlice(desc, localRmaBufferSlice, remoteRmaBufferSlice),
            "[HcclOneSidedConn][FillIpcRmaBufferSlice] Get rmaBufferSlice failed!");
    }
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedConn::FillRdmaRmaBufferSlice(const HcclOneSideOpDesc* desc,
    RmaBufferSlice& localRmaBufferSlice, RmaBufferSlice& remoteRmaBufferSlice)
{
    u32 unitSize;
    CHK_RET(SalGetDataTypeSize(desc->dataType, unitSize));
    u64 byteSize = desc->count * unitSize;

    auto localKey = BufferKey<uintptr_t, u64>(reinterpret_cast<uintptr_t>(desc->localAddr), byteSize);
    auto remoteKey = BufferKey<uintptr_t, u64>(reinterpret_cast<uintptr_t>(desc->remoteAddr), byteSize);
    auto localBuffer = localRdmaRmaBufferMgr_.Find(localKey);
    CHK_PRT_RET(!localBuffer.first,
        HCCL_ERROR("[HcclOneSidedConn][FillRdmaRmaBufferSlice] Cann't find localBuffer by key {%p, %llu},"
            "count is %llu, dataType is %d", desc->localAddr, byteSize, desc->count,
            static_cast<int>(desc->dataType)), HCCL_E_INTERNAL);
    CHK_PRT_RET(!localBuffer.second->GetAddr(), HCCL_ERROR("[HcclOneSidedConn][FillRdmaRmaBufferSlice]"
        "The addr of local Buffer or remote buffer is nullptr."), HCCL_E_NOT_FOUND);
    CHK_PRT_RET(!localBuffer.second->GetDevAddr(), HCCL_ERROR("[HcclOneSidedConn][FillRdmaRmaBufferSlice]"
        "The dev addr of local Buffer is nullptr."), HCCL_E_NOT_FOUND);

    auto remoteBuffer = remoteRdmaRmaBufferMgr_.Find(remoteKey);
    CHK_PRT_RET(!remoteBuffer.first,
        HCCL_ERROR("[HcclOneSidedConn][FillRdmaRmaBufferSlice] Cann't find remoteBuffer by key {%p, %llu},"
            "count is %llu, dataType is %d", desc->remoteAddr, byteSize, desc->count,
            static_cast<int>(desc->dataType)), HCCL_E_INTERNAL);
    CHK_PRT_RET(!remoteBuffer.second->GetDevAddr(), HCCL_ERROR("[HcclOneSidedConn][FillRdmaRmaBufferSlice]"
        "The dev addr of remote buffer is nullptr."), HCCL_E_NOT_FOUND);
    CHK_PRT_RET(!remoteBuffer.second->GetAddr(), HCCL_ERROR("[HcclOneSidedConn][FillRdmaRmaBufferSlice]"
        "The addr of remote buffer is nullptr."), HCCL_E_NOT_FOUND);
    u64 localDataOffSet = static_cast<u8*>(desc->localAddr) - static_cast<u8*>(localBuffer.second->GetAddr());
    u64 remoteDataOffSet = static_cast<u8*>(desc->remoteAddr) - static_cast<u8*>(remoteBuffer.second->GetAddr());
    localRmaBufferSlice.addr = static_cast<void*>(static_cast<u8*>(localBuffer.second->GetDevAddr()) + localDataOffSet);
    localRmaBufferSlice.len = byteSize;
    localRmaBufferSlice.rmaBuffer = localBuffer.second;
    localRmaBufferSlice.memType = localBuffer.second->GetMemType();
    remoteRmaBufferSlice.addr = static_cast<void*>(static_cast<u8*>(remoteBuffer.second->GetDevAddr()) + remoteDataOffSet);
    remoteRmaBufferSlice.len = byteSize;
    remoteRmaBufferSlice.rmaBuffer = remoteBuffer.second;
        remoteRmaBufferSlice.memType = remoteBuffer.second->GetMemType();
    HCCL_INFO("[HcclOneSidedConn][FillRdmaRmaBufferSlice] Local address before mapping is [%p], after mapping is [%p]."
        "Remote address before mapping is [%p], after mapping is [%p]. Datasize is [%llu].",
        desc->localAddr, localRmaBufferSlice.addr, desc->remoteAddr, remoteRmaBufferSlice.addr, byteSize);
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedConn::FillIpcRmaBufferSlice(const HcclOneSideOpDesc* desc,
    RmaBufferSlice& localRmaBufferSlice, RmaBufferSlice& remoteRmaBufferSlice)
{
    u32 unitSize;
    CHK_RET(SalGetDataTypeSize(desc->dataType, unitSize));
    u64 byteSize = desc->count * unitSize;
    auto localKey = BufferKey<uintptr_t, u64>(reinterpret_cast<uintptr_t>(desc->localAddr), byteSize);
    auto remoteKey = BufferKey<uintptr_t, u64>(reinterpret_cast<uintptr_t>(desc->remoteAddr), byteSize);
    auto localBuffer = localIpcRmaBufferMgr_.Find(localKey);
    CHK_PRT_RET(!localBuffer.first,
        HCCL_ERROR("[HcclOneSidedConn][FillIpcRmaBufferSlice] Cann't find localBuffer by key {%p, %llu},"
            "count is %llu, dataType is %d", desc->localAddr, byteSize, desc->count,
            static_cast<int>(desc->dataType)), HCCL_E_INTERNAL);
    CHK_PRT_RET(!localBuffer.second->GetAddr(), HCCL_ERROR("[HcclOneSidedConn][FillIpcRmaBufferSlice]"
        "The addr of local Buffer or remote buffer is nullptr."), HCCL_E_NOT_FOUND);
    CHK_PRT_RET(!localBuffer.second->GetDevAddr(), HCCL_ERROR("[HcclOneSidedConn][FillIpcRmaBufferSlice]"
        "The dev addr of local Buffer is nullptr."), HCCL_E_NOT_FOUND);
    auto remoteBuffer = remoteIpcRmaBufferMgr_.Find(remoteKey);
    CHK_PRT_RET(!remoteBuffer.first,
        HCCL_ERROR("[HcclOneSidedConn][FillIpcRmaBufferSlice] Cann't find remoteBuffer by key {%p, %llu},"
            "count is %llu, dataType is %d", desc->remoteAddr, byteSize, desc->count,
            static_cast<int>(desc->dataType)), HCCL_E_INTERNAL);
    CHK_PRT_RET(!remoteBuffer.second->GetDevAddr(), HCCL_ERROR("[HcclOneSidedConn][FillIpcRmaBufferSlice]"
        "The dev addr of remote buffer is nullptr."), HCCL_E_NOT_FOUND);
    CHK_PRT_RET(!remoteBuffer.second->GetAddr(), HCCL_ERROR("[HcclOneSidedConn][FillIpcRmaBufferSlice]"
        "The addr of remote buffer is nullptr."), HCCL_E_NOT_FOUND);
    u64 localDataOffSet = static_cast<u8*>(desc->localAddr) - static_cast<u8*>(localBuffer.second->GetAddr());
    u64 remoteDataOffSet = static_cast<u8*>(desc->remoteAddr) - static_cast<u8*>(remoteBuffer.second->GetAddr());
    localRmaBufferSlice.addr = static_cast<void*>(static_cast<u8*>(localBuffer.second->GetDevAddr()) + localDataOffSet);
    localRmaBufferSlice.len = byteSize;
    localRmaBufferSlice.rmaBuffer = localBuffer.second;
    localRmaBufferSlice.memType = localBuffer.second->GetMemType();
    remoteRmaBufferSlice.addr = static_cast<void*>(static_cast<u8*>(remoteBuffer.second->GetDevAddr()) + remoteDataOffSet);
    remoteRmaBufferSlice.len = byteSize;
    remoteRmaBufferSlice.rmaBuffer = remoteBuffer.second;
    remoteRmaBufferSlice.memType = remoteBuffer.second->GetMemType();
    HCCL_INFO("[HcclOneSidedConn][FillIpcRmaBufferSlice] Local address before mapping is [%p], after mapping is [%p]."
        "Remote address before mapping is [%p], after mapping is [%p]. Datasize is [%llu].",
        desc->localAddr, localRmaBufferSlice.addr, desc->remoteAddr, remoteRmaBufferSlice.addr, byteSize);
    return HCCL_SUCCESS;
}

}